# -*- coding:utf-8 -*-

import json
from copy import deepcopy


class JsonLoader(object):
    def __init__(self):
        pass

    def loadFirstImage(self, jsonFileName):
        try:
            with open(jsonFileName) as in_file:
                info = json.load(in_file)
        except Exception as e:
            print("error, open json file {} failed, {}".format(jsonFileName, e))
            raise e

        edges = info['links']
        # 获取结点集合, example:
        # ["492:winlogon", "556:explorer"]
        node_names = list(set(e['source'] for e in edges) | set(e['target'] for e in edges))
        nodes = self.nodesFromEdges(edges, node_names)

        network = {'edges': edges, 'nodes': nodes}
        return network

    def nodesFromEdges(self, edges, node_names):
        nodes = {node_names[i]: {'id': str(i)} for i in range(len(node_names))}

        # get node group
        for edge in edges:
            name = edge['source']
            if 'group' not in nodes[name].keys():
                nodes[name]['group'] = edge['source_group']

            name = edge['target']
            if 'group' not in nodes[name].keys():
                nodes[name]['group'] = edge['target_group']

        return nodes

    def loadNextImage(self, network, jsonFileName):
        try:
            with open(jsonFileName) as in_file:
                info = json.load(in_file)
        except Exception as e:
            print("error, open json file {} failed, {}".format(jsonFileName, e))
            raise e

        edges = deepcopy(network["edges"]) + info["links_add"]

        node_names = list(set(e['source'] for e in edges) | set(e['target'] for e in edges))
        nodes = self.nodesFromEdges(edges, node_names)

        # add nodes
        new_nodes = info['nodes_add']
        for newn in new_nodes:
            nodes[newn['name']]['new'] = True

        # 需要删除节点时
        if info["links_remove"] != []:
            for link in info["links_remove"]:
                # network["edges"].remove(link)
                for edge in edges:
                    if edge["source"] == link["source"] and edge["target"] == link["target"]:
                        edges.remove(edge)

            for link in info["nodes_remove"]:
                del nodes[link["name"]]

        network = {'edges': edges, 'nodes': nodes} 
        return network

    def loadVadImage(self, jsonFileName, pid):
        try:
            with open(jsonFileName) as in_file:
                info = json.load(in_file)
        except Exception as e:
            print("error, open json file {} failed, {}".format(jsonFileName, e))
            raise e

        edges = info['edges'][pid]
        nodes = info['nodes'][pid]
        network = {'edges': edges, 'nodes': nodes}
        return network